package com.shujia.flink.core

import org.apache.flink.api.common.functions.RuntimeContext
import org.apache.flink.api.common.state.{ValueState, ValueStateDescriptor}
import org.apache.flink.configuration.Configuration
import org.apache.flink.streaming.api.functions.KeyedProcessFunction
import org.apache.flink.streaming.api.scala._
import org.apache.flink.util.Collector

object Demo14AvgAge {
  def main(args: Array[String]): Unit = {
    val env: StreamExecutionEnvironment = StreamExecutionEnvironment.getExecutionEnvironment

    /**
     *
     * 实时计算每一个班级的平均年龄
     *
     */
    val studentDS: DataStream[String] = env.socketTextStream("master", 8888)


    //取出班级和学生的年龄
    val clazzAndAge: DataStream[(String, Int)] = studentDS.map(line => {
      val split: Array[String] = line.split(",")
      val clazz: String = split(4)
      val age: Int = split(2).toInt
      (clazz, age)
    })
    //按照班级分组
    val keyByDS: KeyedStream[(String, Int), String] = clazzAndAge.keyBy(_._1)

    //计算平均年龄
    //flink的状态可以再任务算子中使用，map,faltmap,filter process都可以
    val avgAgeDS: DataStream[(String, Double)] = keyByDS
      .process(new KeyedProcessFunction[String, (String, Int), (String, Double)] {


        override def open(parameters: Configuration): Unit = {

          /**
           * 定义两个状态来保存总人数和总的年龄
           *
           */
          val context: RuntimeContext = getRuntimeContext

          //总人数状态的描述对象
          val sumNumDesc = new ValueStateDescriptor[Int]("sumNum", classOf[Int])

          //总的年龄的描述对象
          val sumAgeDesc = new ValueStateDescriptor[Int]("sumAge", classOf[Int])

          sumNumState = context.getState(sumNumDesc)
          sumAgeState = context.getState(sumAgeDesc)

        }

        //保存总的人数的状态
        var sumNumState: ValueState[Int] = _
        //保存总的年龄的状态
        var sumAgeState: ValueState[Int] = _

        override def processElement(kv: (String, Int),
                                    ctx: KeyedProcessFunction[String, (String, Int), (String, Double)]#Context,
                                    out: Collector[(String, Double)]): Unit = {

          val clazz: String = kv._1
          val age: Int = kv._2

          //获取之前的总的人数和总的年龄
          var sumNum: Int = sumNumState.value()

          //人数累加
          sumNum += 1

          //更新状态
          sumNumState.update(sumNum)

          var sumAge: Int = sumAgeState.value()
          //累加
          sumAge += age

          //更新状态
          sumAgeState.update(sumAge)

          //季孙平均年龄
          val avgAge: Double = sumAge.toDouble / sumNum

          //将数据发送到下游
          out.collect((clazz, avgAge))
        }
      })


    avgAgeDS.print()

    env.execute()

  }

}
